package panda.natalia.gothread.suspend.instrument;


import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.util.CheckClassAdapter;
import org.objectweb.asm.util.TraceClassVisitor;

import java.io.*;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.nio.file.StandardOpenOption;
import java.util.Date;
import java.util.WeakHashMap;

/**
 * a simple framework of java fiber
 *
 * @Author: scpanda2007
 * @Date: 2019/11/28 15:54
 *
 * 用于修改 原class suspendable 函数
 *  在里面插入 捕获异常，存储上下文，以及还原上下文的 vm code
 *
 */
public final class GoThreadMethodInstrumentor {

    @SuppressWarnings("WeakerAccess")
    public static final int ASMAPI = Opcodes.ASM7;

    private final static String EXAMINED_CLASS = System.getProperty("gothread.writeInstrumentedClasses");
    private WeakHashMap<ClassLoader, MethodDatabase> dbForClassLoader = new WeakHashMap<>();
    private MethodDatabase boostrapDB;

    private boolean check;
    private final boolean aot;
    private boolean verbose;
    private boolean debug;
    private int logLevelMask;

    public GoThreadMethodInstrumentor() { this(false);}

    public GoThreadMethodInstrumentor(boolean aot){
        this.aot = aot;
        setLogLevelMask();
    }

    private synchronized void setLogLevelMask(){
        logLevelMask = (1 << LogLevel.WARNING.ordinal());
        if(verbose || debug) logLevelMask |= (1 << LogLevel.INFO.ordinal());
        if(debug) logLevelMask |= (1 << LogLevel.DEBUG.ordinal());
    }


    private static byte[] toByteArray(InputStream in) throws IOException {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        copy(in, out);
        return out.toByteArray();
    }

    private static final int BUF_SIZE = 8192;

    private static long copy(InputStream from, OutputStream to) throws IOException{
        byte[] buf = new byte[BUF_SIZE];
        long total = 0;
        while (true) {
            int read = from.read(buf);
            if(read == -1) break; // eof
            to.write(buf, 0, read);
            total += read;
        }
        return total;
    }

    public void log(LogLevel level, String msg, Object ...args){
        System.out.println(String.format(msg, args));
    }

    public void error(String msg, Throwable ex){
        System.out.println(msg + ex.getMessage());
    }

    byte[] instrumentClass(ClassLoader classLoader, String className, InputStream is, boolean forceInstrumentation) throws IOException {
        className = className != null ? className.replace('.', '/') : null;

        byte[] cb = toByteArray(is);

        MethodDatabase db = getMethodDatabase(classLoader);

        if(className != null){
            log(LogLevel.INFO, "TRANSFORM: %s %s", className,
                    (db.getClassEntry(className) != null && db.getClassEntry(className).requiresInstrumentation())? "request" : "");
            // preinstr 即在插入动态修改代码之前的原文件做一次保存
            examine(className, "gothread-1-preinstr", cb);
        } else {
            log(LogLevel.INFO, "Transform: null className");
        }

        final ClassReader r1 = new ClassReader(cb);
        final ClassWriter cw1 = new ClassWriter(r1, 0);
        final LabelSuspendableCallSitesClassVisitor ic1 = new LabelSuspendableCallSitesClassVisitor(cw1, db);
        r1.accept(ic1, 0);
        cb = cw1.toByteArray();

        examine(className, "gothread-2", cb);

        final ClassReader r2 = new ClassReader(cb);
        final ClassWriter cw2 = new DBClassWriter(db, r2);
        final ClassVisitor cv2 = (check && EXAMINED_CLASS == null) ? new CheckClassAdapter(cw2) : cw2;
        final InstrumentClass ic2 = new InstrumentClass(cv2, db, forceInstrumentation);

        try {
            r2.accept(ic2, ClassReader.SKIP_FRAMES);
            cb = cw2.toByteArray();
        } catch (final Exception e){
            if(ic2.hasSuspendableMethods()){
                error("Unable to instrument class " + className, e);
                throw e;
            }else{
                if(!MethodDatabase.isProblematicClass(className))
                    log(LogLevel.DEBUG, "Unable to instrument class " + className);
                return null;
            }
        }

        examine(className, "gothread-4", cb);

        final OffsetClassReader r3 = new OffsetClassReader(cb);
        final ClassWriter cw3 = new ClassWriter(r3, 0);
        final SuspendOffsetsAfterInstrClassVisitor ic3 = new SuspendOffsetsAfterInstrClassVisitor(cw3, db);

        r3.accept(ic3, 0);
        cb = cw3.toByteArray();

        if(EXAMINED_CLASS != null){
            examine(className, "gothread-5-final", cb);
            if(check){
                ClassReader r4 = new ClassReader(cb);
                ClassVisitor cv4 = new CheckClassAdapter(new TraceClassVisitor(null), true);
                r4.accept(cv4, 0);
            }
        }

        return cb;
    }

    private void examine(String className, String suffix, byte[] data){
        if(EXAMINED_CLASS != null && className != null && className.contains(EXAMINED_CLASS)){
            final String filename = className.replace('/', '.') + "-" + new Date().getTime() + '-' + suffix + ".class";
            writeToFile(filename, data);
        }
    }

    private static void writeToFile(String name, byte[] data){
        try(OutputStream os = Files.newOutputStream(Paths.get(name), StandardOpenOption.CREATE_NEW)){
            os.write(data);
        }catch (final IOException e){
            throw new RuntimeException(e);
        }
    }

    private MethodDatabase getMethodDatabase(ClassLoader classLoader) {
        if(classLoader == null){
            if(boostrapDB == null){
                boostrapDB = new MethodDatabase(this, null, new DefaultSuspendableClassifier(null));
            }
            return boostrapDB;
        }

        if(!dbForClassLoader.containsKey(classLoader)){
            MethodDatabase newDb = new MethodDatabase(this, classLoader, new DefaultSuspendableClassifier(classLoader));
            dbForClassLoader.put(classLoader, newDb);
            return newDb;
        }

        return dbForClassLoader.get(classLoader);
    }

}
